package checker

import (
	"kumachan/interpreter/lang/textual/ast"
	"kumachan/interpreter/lang/common/limits"
	"kumachan/interpreter/lang/common/source"
	"kumachan/interpreter/compiler/checker/typsys"
	"kumachan/interpreter/compiler/checker/checked"
)


func liftToEnum(expr *checked.Expr, target *typsys.TypeDef) (*checked.Expr, bool) {
	var def, args, ok = typeGetRefDefArgs(expr.Type)
	if ok {
		var chain, ok = typeDefFindEnumChain(def, target)
		if ok {
			return applyEnumChain(expr, args, chain), true
		} else {
			return nil, false
		}
	} else {
		return nil, false
	}
}

func checkIf(I ast.If) ExprChecker {
	return makeExprChecker(I.Location, func(cc *checkContext) checkResult {
		var I = desugarElseIf(I)
		var loc = I.Location
		return cc.forward(checkIfElse(I.Condition, I.YesBranch, I.NoBranch, loc))
	})
}
func checkIfElse(cond ast.Expr, yes ast.Expr, no ast.Expr, loc source.Location) ExprChecker {
	return makeExprChecker(loc, func(cc *checkContext) checkResult {
		var result_t, _ = cc.getExpectedCertainOrInferred(cc.expected)
		var cond_expr, err1 = cc.checkChildExpr(cc.getType(coreBool), cond)
		if err1 != nil { return cc.propagate(err1) }
		var yes_expr, err2 = cc.checkChildExpr(result_t, yes)
		if err2 != nil { return cc.propagate(err2) }
		if result_t == nil { result_t = yes_expr.Type }
		var no_expr, err3 = cc.checkChildExpr(result_t, no)
		if err3 != nil { return cc.propagate(err3) }
		var branches = make(map[uint] checked.Branch)
		branches[coreYesIndex] = checked.Branch { Lambda:
			checked.Lambda {
				In:  checked.ProductPatternInfo(nil),
				Out: yes_expr,
			} }
		branches[coreNoIndex] = checked.Branch { Lambda:
			checked.Lambda {
				In:  checked.ProductPatternInfo(nil),
				Out: no_expr,
			} }
		return cc.assign(
			result_t,
			checked.Switch {
				Argument:   cond_expr,
				HasDefault: false,
				Branches:   branches,
			})
	})
}

func checkSwitch(S ast.Switch) ExprChecker {
	return makeExprChecker(S.Location, func(cc *checkContext) checkResult {
		var result_t, _ = cc.getExpectedCertainOrInferred(cc.expected)
		var arg, err = cc.checkChildExpr(nil, S.Argument)
		if err != nil { return cc.propagate(err) }
		var enum, _, type_args, is_enum = cc.unboxEnum(arg)
		if !(is_enum) {
			return cc.error(
				E_SwitchOnNonEnum {})
		}
		var case_index_map = make(map[string] uint)
		for i, case_def := range enum.CaseTypes {
			case_index_map[case_def.Name.ItemName] = uint(i)
		}
		var branches = make(map[uint] checked.Branch)
		var default_branch checked.Branch
		var has_default = false
		for _, B := range S.Branches {
			if len(B.Cases) == 0 {
				// default branch
				if has_default {
					return cc.propagate(source.MakeError(B.Location,
						E_DuplicateCase {}))
				}
				has_default = true
				var expr, err = cc.checkChildExpr(result_t, B.Expr)
				if err != nil { return cc.propagate(err) }
				if result_t == nil { result_t = expr.Type }
				var in = [] checked.ProductPatternItemInfo { {
					Binding: nil,
					Index1:  0,
				} }
				default_branch = checked.Branch { Lambda:
					checked.Lambda {
						In:  in,
						Out: expr,
					} }
				continue
			}
			for _, case_node := range B.Cases {
				var case_str = ast.Id2String(case_node)
				var case_index, exists = case_index_map[case_str]
				if !(exists) {
					return cc.propagate(source.MakeError(case_node.Location,
						E_NoSuchCase {
							Case: case_str,
						}))
				}
				var _, occurred = branches[case_index]
				if occurred {
					return cc.propagate(source.MakeError(case_node.Location,
						E_DuplicateCase {}))
				}
				var loc = B.Location
				var out = B.Expr
				var in = B.Pattern
				var in_t = &typsys.NestedType { Content:
					typsys.Ref {
						Def:  enum.CaseTypes[case_index],
						Args: type_args,
					} }
				var branch_checker = checkBranch(in_t, in, out, loc)
				var expr, err = cc.checkChild(result_t, branch_checker)
				if err != nil { return cc.propagate(err) }
				if result_t == nil { result_t = expr.Type }
				branches[case_index] = expr.Content.(checked.Branch)
			}
		}
		if !(has_default) {
			if len(branches) != len(enum.CaseTypes) {
				return cc.error(
					E_CondNotExhaustive {})
			}
		}
		return cc.assign(
			result_t,
			checked.Switch {
				Argument:   arg,
				HasDefault: has_default,
				Default:    default_branch,
				Branches:   branches,
			})
	})
}

func checkSelect(S ast.Select) ExprChecker {
	return makeExprChecker(S.Location, func(cc *checkContext) checkResult {
		var result_t, _ = cc.getExpectedCertainOrInferred(cc.expected)
		var arity = uint(len(S.Arguments))
		if arity > limits.MaxSelectArity {
			return cc.error(
				E_SelectTooManyArguments {})
		}
		if arity < 2 {
			return cc.error(
				E_SelectInsufficientArguments {})
		}
		type Arg struct {
			Expr      *checked.Expr
			Enum      *typsys.Enum
			Cases     map[string] uint
			TypeArgs  [] typsys.Type
		}
		var args = make([] Arg, arity)
		for i, arg_node := range S.Arguments {
			var expr, err = cc.checkChildExpr(nil, arg_node)
			if err != nil { return cc.propagate(err) }
			var enum, _, type_args, is_enum = cc.unboxEnum(expr)
			if !(is_enum) {
				return cc.propagate(source.MakeError(arg_node.Location,
					E_SwitchOnNonEnum {}))
			}
			var case_index_map = make(map[string] uint)
			for i, case_def := range enum.CaseTypes {
				case_index_map[case_def.Name.ItemName] = uint(i)
			}
			args[i] = Arg {
				Expr:     expr,
				Enum:     enum,
				Cases:    case_index_map,
				TypeArgs: type_args,
			}
		}
		var branches = make(map[checked.SelectBranchKey] checked.Branch)
		var default_branch checked.Branch
		var has_default = false
		for _, B := range S.Branches {
			if len(B.Cases) == 0 {
				// default branch
				if has_default {
					return cc.propagate(source.MakeError(B.Location,
						E_DuplicateCase {}))
				}
				has_default = true
				var expr, err = cc.checkChildExpr(result_t, B.Expr)
				if err != nil { return cc.propagate(err) }
				if result_t == nil { result_t = expr.Type }
				var in = [] checked.ProductPatternItemInfo { {
					Binding: nil,
					Index1:  0,
				} }
				default_branch = checked.Branch { Lambda:
					checked.Lambda {
						In:  in,
						Out: expr,
					} }
				continue
			}
			for _, case_node := range B.Cases {
				var sub_case_nodes = case_node.Components
				if uint(len(sub_case_nodes)) != arity {
					return cc.error(
						E_TupleSizeNotMatching {
							Required: arity,
							Given:    uint(len(sub_case_nodes)),
						})
				}
				var key checked.SelectBranchKey
				var in_elements = make([] typsys.Type, arity)
				for i, sub_case_node := range sub_case_nodes {
					var arg = &(args[i])
					var sub_case_str = ast.Id2String(sub_case_node)
					var sub_case_index, exists = arg.Cases[sub_case_str]
					if !(exists) {
						return cc.propagate(source.MakeError(sub_case_node.Location,
							E_NoSuchCase {
								Case: sub_case_str,
							}))
					}
					key[i] = sub_case_index
					in_elements[i] = &typsys.NestedType { Content:
						typsys.Ref {
							Def:  arg.Enum.CaseTypes[sub_case_index],
							Args: arg.TypeArgs,
						} }
				}
				var _, occurred = branches[key]
				if occurred {
					return cc.propagate(source.MakeError(case_node.Location,
						E_DuplicateCase {}))
				}
				var in_t = &typsys.NestedType { Content:
					typsys.Tuple { Elements:
						in_elements,
					} }
				var in = B.Pattern
				var out = B.Expr
				var loc = B.Location
				var branch_checker = checkBranch(in_t, in, out, loc)
				var expr, err = cc.checkChild(result_t, branch_checker)
				if err != nil { return cc.propagate(err) }
				if result_t == nil { result_t = expr.Type }
				branches[key] = expr.Content.(checked.Branch)
			}
		}
		for U := range branches {
			for V := range branches {
				if U != V {
					var overlap = true
					for i := uint(0); i < arity; i += 1 {
						var u = U[i]
						var v = V[i]
						var overlap_i = (u == v)
						overlap = (overlap && overlap_i)
					}
					if overlap {
						return cc.error(
							E_SelectOverlappingBranches {})
					}
				}
			}
		}
		if !(has_default) {
			var total = 1
			for _, arg := range args {
				total = (total * len(arg.Enum.CaseTypes))
			}
			if len(branches) != total {
				return cc.error(
					E_CondNotExhaustive {})
			}
		}
		var args_expr = make([] *checked.Expr, len(args))
		for i, arg := range args {
			args_expr[i] = arg.Expr
		}
		return cc.assign(
			result_t,
			checked.Select {
				Arguments:  args_expr,
				HasDefault: has_default,
				Default:    default_branch,
				Branches:   branches,
			})
	})
}

func getEnum(t typsys.Type) (*typsys.Enum, *typsys.TypeDef, ([] typsys.Type), bool) {
	var nested, is_nested = t.(*typsys.NestedType)
	if !(is_nested) { return nil, nil, nil, false }
	var ref, is_ref = nested.Content.(typsys.Ref)
	if !(is_ref) { return nil, nil, nil, false }
	var enum, is_enum = ref.Def.Content.(*typsys.Enum)
	if !(is_enum) { return nil, nil, nil, false }
	return enum, ref.Def, ref.Args, true
}
func unboxEnum(t typsys.Type, mod string) (*typsys.Enum, *typsys.TypeDef, ([] typsys.Type), bool) {
	var enum, def, args, is_enum = getEnum(t)
	if is_enum {
		return enum, def, args, true
	} else {
		var inner, _, exists = typsys.Unbox(t, mod)
		if exists {
			return unboxEnum(inner, mod)
		} else {
			return nil, nil, nil, false
		}
	}
}
func applyEnumChain(expr *checked.Expr, args ([] typsys.Type), chain ([] typsys.CaseInfo)) *checked.Expr {
	for _, info := range chain {
		var enum = info.Enum
		var index = info.CaseIndex
		if len(args) != len(enum.Parameters) {
			panic("something went wrong")
		}
		var t = &typsys.NestedType { Content:
			typsys.Ref {
				Def:  enum,
				Args: args,
			} }
		expr = &checked.Expr {
			Type:    t,
			Info:    expr.Info,
			Content: checked.Enum {
				Index: index,
				Value: expr,
			},
		}
	}
	return expr
}
func desugarElseIf(raw ast.If) ast.If {
	var no_branch = raw.NoBranch
	var elifs = raw.ElIfs
	for i, _ := range elifs {
		var elif = elifs[len(elifs)-1-i]
		var t = ast.If {
			Node:      elif.Node,
			Condition: elif.Condition,
			YesBranch: elif.YesBranch,
			NoBranch:  no_branch,
			ElIfs:     nil,
		}
		no_branch = ast.WrapTermAsExpr(ast.VariousTerm {
			Node: t.Node,
			Term: t,
		})
	}
	return ast.If {
		Node:      raw.Node,
		Condition: raw.Condition,
		YesBranch: raw.YesBranch,
		NoBranch:  no_branch,
		ElIfs:     nil,
	}
}


